package top.yudoge.dogeeditor.llm.rewrite

import com.google.gson.Gson
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.app.utils.HistoryMessageHolder
import top.yudoge.dogeeditor.app.utils.Message
import top.yudoge.dogeeditor.commons.utils.convertList
import top.yudoge.dogeeditor.commons.utils.jsonToBean
import top.yudoge.dogeeditor.commons.utils.toJsonStr
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.embedding.EmbeddingStoreHolder
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.objects.mro.TableMRO
import top.yudoge.dogeeditor.llm.task.TableEmbeddingTask
import top.yudoge.dogeeditor.llm.util.VectorUtils
import top.yudoge.dogeeditor.llm.util.getRelatedTableNames

@Component
class GetPossibleTableChainNode : TaskChainNode {

    @Autowired
    private lateinit var tableEmbeddingTask: TableEmbeddingTask

    @Autowired
    private lateinit var embeddingStoreHolder: EmbeddingStoreHolder

    @Autowired
    private lateinit var dataSourceMetaDataApi: DataSourceMetaDataApi

    @Autowired
    private lateinit var sessionMemory: SessionMemory

    override fun process(ctx: ChainContext) {
        var tableNamesFromRAG = ragTables(ctx)
        filterTablesByModel(ctx, tableNamesFromRAG)
    }

    private fun filterTablesByModel(ctx: ChainContext, tablesFromRag: List<String>) {
        sendReasoningCallback(ctx, "Exact matching")

        val historyTables = safeReadJsonArray(sessionMemory.get(ctx.query.sessionId, SQLGenerate2ChainNode.HIS_TABLES_KEY) ?: "[]", String::class.java)

        val allTableNames = mutableSetOf<String>().apply {
            addAll(tablesFromRag)
            addAll(historyTables)
            addAll(getChainStore(ctx, UserInputAnalyseChainNode.TABLE_NAME_KEY) as? List<String> ?: ArrayList())
        }

        if (allTableNames.isEmpty()) {
            putChainStore(ctx, POSSIBLE_TABLES, listOf<String>())
            emitNext(ctx)
            return
        }

        val allTableMROs = mutableListOf<TableMRO>()

        for (tableName in allTableNames) {
            var columns =
                dataSourceMetaDataApi.getColumns(ctx.query.dataSourceIdentifier, ctx.query.databaseName, ctx.query.schemaName, tableName)
            if (!columns.isNullOrEmpty()) {
                var tableMRO = TableMRO(tableName, columns.convertList())
                allTableMROs.add(tableMRO)
            }
        }

        if (allTableMROs.isEmpty()) {
            putChainStore(ctx, POSSIBLE_TABLES, listOf<String>())
            emitNext(ctx)
            return
        }

        var message = mutableListOf<Message>().apply {
            addAll(listChatHistory(ctx).filter { it.type == "user" })
            add(Message(type = "user", content = ctx.query.userTask))
        }.toJsonStr()

        var prompt = loadPrompt(
            "/guess_table.template", mapOf(
                "user_task" to message,
                "tables" to allTableMROs.toJsonStr()
            )
        )

        var response = ctx.chatModel.chat(prompt)

        println(response)

        val result = safeReadJson(response).jsonToBean<Result>()

        if (result.valid) {
            val tableNames =result.tableNames ?: listOf()
            if (tableNames.isNotEmpty()) {
                sendText(ctx, "Tables: ")
                sendText(ctx, tableNames.map { "- `${it}`"}.joinToString("\n"))
                sendText(ctx, "")
            }
            println("Filtered table names $tableNames")
            putChainStore(ctx, POSSIBLE_TABLES, tableNames)
            emitNext(ctx)
        } else {
            ctx.callback.token(result.reason ?: "")
            ctx.callback.complete()
        }

    }


    private fun ragTables(ctx: ChainContext): List<String> {

        sendReasoningCallback(ctx, "RAG Tables")

        // todo 此处可能使用用户历史上所有的提问做RAG
        var query = ctx.query

        var userTask = ctx.query.userTask
        var userTaskEmbedding = tableEmbeddingTask.embedding(userTask)

        var embeddingStore = embeddingStoreHolder.getStore(query.dataSourceIdentifier, query.databaseName, query.schemaName)
        var relatedTableNames = embeddingStore.getRelatedTableNames(userTaskEmbedding, 20, embeddingStoreHolder)

        println("Tables from RAG = ${relatedTableNames}")

        return relatedTableNames
    }

    private fun addChatHistory(ctx: ChainContext, message: Message) {
        var historyHolder = getHistoryHolder(ctx)
        historyHolder.add(message)
        sessionMemory.put(ctx.query.sessionId, HISTORY_HOLDER_KEY, Gson().toJson(historyHolder))
    }

    private fun listChatHistory(ctx: ChainContext): List<Message> {
        return getHistoryHolder(ctx).history()
    }

    private fun getHistoryHolder(ctx: ChainContext): HistoryMessageHolder {
        return sessionMemory.get(ctx.query.sessionId, HISTORY_HOLDER_KEY)?.let {
            Gson().fromJson(it, HistoryMessageHolder::class.java)
        } ?: HistoryMessageHolder()
    }

    companion object {
        const val POSSIBLE_TABLES = "possible_tables"
        private const val HISTORY_HOLDER_KEY = "message:history"
    }

    class Result {
        var valid: Boolean = false
        var reason: String? = null
        var tableNames: List<String>? = null
    }


}